# server.py
import models
import torch
import torch.utils.data
import torch.nn.functional


class Server(object):
	
	def __init__(self, conf, eval_dataset):
	
		self.conf = conf 
		
		self.global_model = models.get_model(self.conf["model_name"])  # 根据配置文件的“model——name”键得到参数初始值的网络模型
		
		self.eval_loader = torch.utils.data.DataLoader(eval_dataset, batch_size=self.conf["batch_size"], shuffle=True)
		
	# 模型聚合，全局模型的更新
	def model_aggregate(self, weight_accumulator):
		for name, data in self.global_model.state_dict().items():
			
			update_per_layer = weight_accumulator[name] * self.conf["lambda"]
			
			if data.type() != update_per_layer.type():
				data.add_(update_per_layer.to(torch.int64))  # 转为整型64位
			else:
				data.add_(update_per_layer)

	# 模型评估
	def model_eval(self):
		self.global_model.eval()
		
		total_loss = 0.0
		correct = 0
		dataset_size = 0
		for batch_id, batch in enumerate(self.eval_loader):  # batch_id就为enumerate()遍历集合所返回的批量序号
			data, target = batch  # 得到数据集和标签
			dataset_size += data.size()[0]  # data.size()=[batch,通道数,32,32]、target.size()=[batch]
			
			if torch.cuda.is_available():  # 将评估数据放到GPU环境
				data = data.cuda()
				target = target.cuda()

			output = self.global_model(data)
			
			total_loss += torch.nn.functional.cross_entropy(output, target, reduction='sum').item()  # sum up batch loss
			# .data意即将变量的tensor取出来
			# 因为tensor包含data和grad，分别放置数据和计算的梯度
			pred = output.data.max(1)[1]  # get the index of the max log-probability
			# 按照从左往右的 第一维 取出最大值的索引 torch.max()
			correct += pred.eq(target.data.view_as(pred)).cpu().sum().item()
			# torch.view_as(tensor)即将调用函数的变量，转变为同参数tensor同样的形状
			# torch.eq()对两个张量tensor进行逐元素比较，如果相等则返回True，否则返回False。True和False作运算时可以作1、0使用
			# .cpu()这一步将预测结果放到cpu上，利用电脑内存存储列表值。从而避免测试过程中爆显存。
			# .sum()是将我们一个批量的预测值求和，便于累加到correct变量中。
			# .item()取出 单元素张量的元素值 并返回该值，保持原元素类型不变。

		acc = 100.0 * (float(correct) / float(dataset_size))
		total_l = total_loss / dataset_size

		return acc, total_l
